"""
调用基座大模型
creat at:
"""
from openai import OpenAI, AsyncOpenAI


class LLM:
    def __init__(self,
                 model_name: str = None,
                 api_key: str = None,
                 base_url: str = None,
                 temperature: float = 0,
                 max_tokens: int = 1024,
                 top_p: float = 1.0,
                 max_retries: int = 3,
                 delay: float = 0.1):
        """
        Args:
            model_name: 模型名称，如空则从环境变量 DEFAULT_LLM 获取
            api_key: 密钥，如空则从环境变量 LLM_API_KEY 中获取
            base_url: API URL, 如空则从环境变量 LLM_BASE_URL 中获取
            temperature: 温度，0~1
            max_tokens: 生成内容最大长度
            top_p: 生成采样的控制参数
        """
        self.model_name = model_name
        self.api_key = api_key
        self.base_url = base_url

        self.temperature = temperature
        self.max_tokens = max_tokens
        self.top_p = top_p

        self.client = OpenAI(
            api_key=self.api_key,
            base_url=self.base_url,
        )

    def chat(self, prompt: str):
        """
        chat 与大模型进行对话
        :param prompt:
        :return:
        """

        content = prompt

        try:
            # 检查 prompt 类型
            if isinstance(prompt, str):
                messages = [{'role': 'user', 'content': content}]
            elif isinstance(prompt, list):
                messages = prompt
            else:
                return ''

            completion = self.client.chat.completions.create(
                model=self.model_name,
                messages=messages,
                temperature=self.temperature,
                max_tokens=self.max_tokens,
                top_p=self.top_p,
                timeout=9999
            )

            if completion and completion.choices:
                content = completion.choices[0].message.content.strip()
                return content

        except Exception as e:
            print(e)
            return ''


if __name__ == '__main__':
    api_key: str = 'sk-68ac5f5ccf3540ba834deeeaecb48987'
    base_url: str = "https://dashscope.aliyuncs.com/compatible-mode/v1"
    model_name: str = "deepseek-v3"

    llm_t0 = LLM(api_key=api_key,
                 base_url=base_url,
                 model_name=model_name,
                 temperature=0)

    llm_t1 = LLM(api_key=api_key,
                 base_url=base_url,
                 model_name=model_name,
                 temperature=0)

    llm_t2 = LLM(api_key=api_key,
                 base_url=base_url,
                 model_name=model_name,
                 temperature=1)

    resp0 = llm_t0.chat(prompt='用50字解释地球为什么是圆的？')
    resp1 = llm_t1.chat(prompt='用50字解释地球为什么是圆的？')
    resp2 = llm_t2.chat(prompt='用50字解释地球为什么是圆的？')
    print(resp0)
    print(resp1)
    print(resp2)
